#include <iostream>
#include <vector>
#include "diversity.hpp"
#include "time_memory.hpp"
#include "omp.h"
#include <set>
#include <iterator>
#include <algorithm>

using namespace std;
using namespace Diversity;

int main(int argc, char **argv) {
    if (argc != 4) {
        std::cout << argv[0] << " method_path n_classifier n_cluster" << std::endl;
        exit(-1);
    }
    printf("method_path %s\nn_classifier %s\nn_cluster %s\n", argv[1], argv[2], argv[3]);
    char *base_path = argv[1];
    int n_classifier = (int) atoi(argv[2]);
    int n_cluster = (int) atoi(argv[3]);
    int report_every_ = 10;
    if (n_classifier <= 1) {
        printf("number of classifier is less than 1, can not get the result\n");
    }

    vector<vector<set<int>>> label_map_l;
    for (int i = 0; i < n_classifier; i++) {
        char path[256];
        sprintf(path, "%s/Classifier_%d/dataset_partition/label.txt", base_path, i);
        printf("%s\n", path);
        vector<int> partition_l = load_partition_txt(path);
        vector<set<int>> label_map = get_label_map(partition_l, n_cluster);
        label_map_l.push_back(label_map);
    }
    printf("get label map\n");

    vector<pair<int, int>> pair_vector;
    for (int i = 0; i < n_classifier; i++) {
        for (int j = 0; j < n_classifier; j++) {
            if(i != j){
                pair_vector.push_back(make_pair(i, j));
            }
        }
    }
    double diversity = 0;
    int n_iter = n_classifier * (n_classifier - 1);
#pragma omp parallel for
    for (int idx = 0; idx < n_iter; idx++) {
        pair<int, int> pair_idx = pair_vector[idx];
        int i = pair_idx.first;
        int j = pair_idx.second;
        vector<set<int>> i_label_map = label_map_l[i];
        vector<set<int>> j_label_map = label_map_l[j];
        double total_jaccard = 0;

        TimeRecord record;
        int thread_id = omp_get_thread_num();

        for (int i_1 = 0; i_1 < n_cluster; i_1++) {
            double max_jaccard = 0;
            for (int j_1 = 0; j_1 < n_cluster; j_1++) {
                double intersection = intersection_size(i_label_map[i_1], j_label_map[j_1]);
                double union_ = union_size(i_label_map[i_1], j_label_map[j_1]);
                double tmp_jaccard = intersection / union_;
                max_jaccard = max(max_jaccard, tmp_jaccard);
            }
            total_jaccard += max_jaccard;

            if (i_1 % report_every_ == 0) {
                std::cout << i_1 / (0.01 * n_cluster) << " %, "
                          << 1e-6 * record.get_elapsed_time_micro() << " s/iter, thread num " << thread_id << "\n";
                record.reset();
            }

        }


#pragma omp critical
        {
            diversity += total_jaccard;
        }


    }
    diversity = diversity / n_iter / n_cluster;
    printf("diversity of %s: %.3f\n", base_path, diversity);

    return 0;
}